;; Tests of the LUP decomposition functions
;; Copyright (C) 2024  Rocks Mazama
;;
;; This program is free software: you can redistribute it and/or modify
;; it under the terms of the GNU General Public License as published by
;; the Free Software Foundation, either version 3 of the License, or
;; (at your option) any later version.
;;
;; This program is distributed in the hope that it will be useful,
;; but WITHOUT ANY WARRANTY; without even the implied warranty of
;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
;; GNU General Public License for more details.
;;
;; You should have received a copy of the GNU General Public License
;; along with this program.  If not, see <https://www.gnu.org/licenses/>.

(ns mazama.linalg.LUP-test
  (:require
   [clojure.test :refer [deftest is testing]]
   [mazama.linalg.LUP :refer [LUP]]))

(defn columns
  "Return a sequence of columns, each column being a sequence of elements."
  ([{:keys [ncols data]}]
   (columns ncols data))
  ([ncols data]
   (letfn [(f [i data]
             (if (< i ncols)
               (lazy-seq (cons (take-nth ncols data) (f (inc i) (rest data))))
               []))]
     (f 0 (seq data)))))

(defn rows
  "Return a sequence of rows, each row being a sequence of elements."
  ([{:keys [ncols data]}]
   (rows ncols data))
  ([ncols data]
   (partition ncols data)))

(defn dot-product
  [u v]
  (reduce + 0 (mapv * u v)))

(defn- mmul
  "Multiply two matrices, `A` and `B`."
  [A B]
  {:ncols (:ncols B)
   :data (for [row (rows A) col (columns B)] (dot-product row col))})

(defn- matrix-index
  "Return the zero-based indexes of the `i`th element of a matrix with `ncols` columns."
  [ncols i]
  (vector (quot i ncols) (rem i ncols)))

(defn- every-indexed?
  ""
  [f coll]
  (loop [s (seq coll)
         i 0]
    (if s
      (and (f i (first s)) (recur (next s) (inc i)))
      true)))

(defn- unit-lower? [[i j] x]
  (or (> i j)
      (and (= i j) (= x 1))
      (zero? x)))

(defn- unit-lower-triangular?
  "Is the matrix unit lower-triangular?"
  [{:keys [ncols data]}]
  (every-indexed? #(unit-lower? (matrix-index ncols %1) %2) data))

(defn- upper? [[i j] x]
  (or (<= i j) (zero? x)))

(defn- upper-triangular?
  "Is the matrix upper-triangular?"
  [{:keys [ncols data]}]
  (every-indexed? #(upper? (matrix-index ncols %1) %2) data))

(defn- one-and-zeros?
  "Does the `coll` contain exactly one 1 and the rest are all 0?"
  [coll]
  (let [freq (frequencies coll)
        ones (or (freq 1) 0)
        zeros (or (freq 0) 0)]
    (and (== ones 1)
         (== zeros (dec (count coll))))))

(defn- permutation-matrix?
  "Is `m` a permutation matrix?"
  [m]
  (let [f one-and-zeros?]
    (and (every? f (rows m))
         (every? f (columns m)))))

(defn- equal-matrices? [A B]
  (and (== (:ncols A) (:ncols B))
       (or (= (:data A) (:data B))
           (->> (map == (:data A) (:data B))
                (reduce #(and %1 %2) true)))))

(defn- dimension?
  "Does matrix `A` have dimensions `m`x`n`?"
  [A m n]
  (let [ncols (:ncols A)
        nrows (if (zero? ncols) 0 (/ (count (:data A)) ncols))]
    (and (== m nrows) (== n ncols))))

(defn- checkLUP [A]
  (let [{:keys [L U P]} (LUP A)]
    (let [n (:ncols A)
          m (if (zero? n) 0 (/ (count (:data A)) n))]
      (is (dimension? L m (min m n)))
      (is (dimension? U (min m n) n))
      (is (dimension? P m m)))
    ;; Tests for all these predicates are further below
    (is (unit-lower-triangular? L))
    (is (upper-triangular? U))
    (is (permutation-matrix? P))
    (is (equal-matrices? (mmul P A) (mmul L U)))))

(deftest small-test
  (testing "0x0"
    (checkLUP {:ncols 0 :data []}))
  (testing "1x1"
    (checkLUP {:ncols 1 :data [-37]})))

(deftest wikipedia-test
  (testing "https://en.wikipedia.org/wiki/LU_decomposition"
    (checkLUP {:ncols 3 :data [0 5 22/3
                               4 2 1
                               2 7 9]})))

(deftest CLRS-test
  (testing "LUP-Solve"
    (checkLUP {:ncols 3 :data [1 2 0
                               3 4 4
                               5 6 3]}))
  (testing "Figure 28.1"
    (checkLUP {:ncols 4 :data [2  3  1  5
                               6 13  5 19
                               2 19 10 23
                               4 10 11 31]}))
  (testing "Figure 28.2"
    (checkLUP {:ncols 4 :data [ 2  0   2  0.6
                                3  3   4 -2
                                5  5   4  2
                               -1 -2 3.4 -1]}))
  (testing "Figure 28.2 Rationals"
    (checkLUP {:ncols 4 :data [ 2  0     2  6/10
                                3  3     4 -2
                                5  5     4  2
                               -1 -2 34/10 -1]})))

(deftest non-square-test
  (testing "LA with Applications by Nicholson, Example 2.7.2"
    (checkLUP {:ncols 5 :data [0  2 -6 -2  4
                               0 -1  3  3  2
                               0 -1  3  7 10]}))
  (testing "4x3"
    (checkLUP {:ncols 3 :data [-1 -2 -3
                               -7 10 37
                                8  7  6
                               19  0 -2]}))
  (testing "1xwide"
    (checkLUP {:ncols 7 :data [77 66 55 44 33 22 11]}))
  (testing "1xwide zeros"
    (checkLUP {:ncols 6 :data [0 0 0 0 0 0]}))
  (testing "tallx1"
    (checkLUP {:ncols 1 :data [0 2 4 6 7 5 3 1]}))
  (testing "tallx1 zeros"
    (checkLUP {:ncols 1 :data [0 0 0 0 0]})))

(deftest data-as-seq-test
  (testing "range"
    (checkLUP {:ncols 17 :data (range 0 (* 17 33))}))
  (testing "vector seq"
    (checkLUP {:ncols 3 :data (seq [6 5 4 3 2 1])})))

(deftest unit-lower-triangular?-test
  (testing "small"
    (is (unit-lower-triangular? {:ncols 1 :data [1]})))
  (testing "square true"
    (is (unit-lower-triangular? {:ncols 4 :data [1 0 0 0
                                                 2 1 0 0
                                                 3 4 1 0
                                                 5 6 8 1]})))
  (testing "square false"
    (is (not (unit-lower-triangular? {:ncols 3 :data [1 0 0
                                                      2 2 0
                                                      3 4 1]}))))
  (testing "wide true"
    (is (unit-lower-triangular? {:ncols 3 :data [1 0 0
                                                 2 1 0]})))
  (testing "wide false"
    (is (not (unit-lower-triangular? {:ncols 4 :data [1 0 0  0
                                                      3 1 0 -1
                                                      2 1 0  1]}))))
  (testing "tall true"
    (is (unit-lower-triangular? {:ncols 3 :data [1 0 0
                                                 2 1 0
                                                 3 0 1
                                                 4 8 7]})))
  (testing "tall false"
    (is (not (unit-lower-triangular? {:ncols 2 :data [2 1
                                                      1 0
                                                      0 0
                                                      0 0]})))))

(deftest upper-triangular?-test
  (testing "small"
    (is (upper-triangular? {:ncols 1 :data [1]})))
  (testing "square true"
    (is (upper-triangular? {:ncols 4 :data [3 9 8 5
                                            0 4 0 1
                                            0 0 5 0
                                            0 0 0 6]})))
  (testing "square false"
    (is (not (upper-triangular? {:ncols 3 :data [1 0 0
                                                 2 2 0
                                                 3 4 1]}))))
  (testing "non-square wide true"
    (is (upper-triangular? {:ncols 3 :data [0 0 0
                                            0 0 0]})))
  (testing "non-square wide false"
    (is (not (upper-triangular? {:ncols 4 :data [1 0 0  0
                                                 0 2 0 -1
                                                 0 1 0  1]}))))
  (testing "non-square tall true"
    (is (upper-triangular? {:ncols 3 :data [1 2 3
                                            0 1 0
                                            0 0 4
                                            0 0 0]})))
  (testing "non-square tall false"
    (is (not (upper-triangular? {:ncols 2 :data [2 1
                                                 1 0
                                                 0 0
                                                 0 0]})))))

(deftest equal-matrices?-test
  (testing "true"
    (is (equal-matrices? {:ncols 4, :data [5 5  4.0M 2.0 2  0  2.0 0.6 -1  -2  3.4 -1.0 3   3   4.0 -2.0]}
                         {:ncols 4, :data [5 5N 4    2.0 2N 0N 2.0 3/5 -1N -2N 3.4 -1.0 3.0 3.0 4.0 -2]})))
  (testing "false ncols"
    (is (not (equal-matrices? {:ncols 3, :data [1 2 3 4 5 6 7 8 9 10 11 12]}
                              {:ncols 4, :data [1 2 3 4 5 6 7 8 9 10 11 12]}))))
  (testing "false data"
    (is (not (equal-matrices? {:ncols 2 :data [1 2 4 3 5 6]}
                              {:ncols 2 :data [1 2 3 3 5 6]})))))

(deftest permutation-matrix?-test
  (testing "false too many ones"
    (is (not (permutation-matrix? {:ncols 3 :data [0 1 0
                                                   1 0 1
                                                   0 0 1]}))))
  (testing "false something other than 0's and 1"
    (is (not (permutation-matrix? {:ncols 4 :data [0 0 0 1
                                                   0 1 0 2
                                                   1 0 0 0
                                                   0 0 1 0]}))))
  (testing "true"
    (is (permutation-matrix? {:ncols 4 :data [0 0 0 1
                                              0 1 0 0
                                              1 0 0 0
                                              0 0 1 0]}))))

(deftest mmul-test
  (testing "square"
    (is (equal-matrices? (mmul {:ncols 3 :data [1 2 3
                                                3 2 1
                                                1 2 3]}
                               {:ncols 3 :data [4 5 6
                                                6 5 4
                                                4 6 5]})
                         {:ncols 3 :data [28 33 29
                                          28 31 31
                                          28 33 29]})))
  (testing "non-square"
    (is (equal-matrices? (mmul {:ncols 2 :data [2 3
                                                2 1
                                                2 3]}
                               {:ncols 3 :data [4 5 6
                                                6 5 4]})
                         {:ncols 3 :data [26 25 24
                                          14 15 16
                                          26 25 24]}))))
